# -*- coding: utf-8 -*-
# Author  : liyanpeng
# Email   : yanpeng.li@cumt.edu.cn
# Datetime: 2024/7/14 13:04
# Filename: cmteb.py
from mteb import MTEB
from C_MTEB import ChineseTaskList, load_retrieval_data
from sentence_transformers import SentenceTransformer
from datasets import load_dataset
import os
from prettytable import PrettyTable


def print_table(task_names, scores):
    tb = PrettyTable()
    tb.field_names = task_names
    tb.add_row(scores)
    print(tb)


# os.environ['HF_DATASETS_CACHE'] = data_path
# os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'


if __name__ == '__main__':
    data_path = '/root/data3/liyanpeng/hf_data'
    model_name = "/root/data3/emb_models/sensenova/piccolo-base-zh"

    model = SentenceTransformer(model_name)
    evaluation = MTEB(tasks=ChineseTaskList)
    for task in evaluation.tasks:
        if task.description.get('type') == 'Retrieval':
            task.corpus, task.queries, task.relevant_docs = load_retrieval_data(
                os.path.join(data_path, task.description["hf_hub_name"]),
                task.description['eval_splits'])
        else:
            dataset = load_dataset(
                path=os.path.join(data_path, task.description["hf_hub_name"]),
                revision=task.description.get("revision", None),
                # task_eval_splits=task.description.get("eval_splits", [])
            )
            task.dataset = dataset
        task.data_loaded = True
    results = evaluation.run(model, output_folder=f"zh_results/piccolo-base-zh")
    print(results)
